package com.ibm.cps.kafka;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Properties;

import org.apache.log4j.Logger;

import kafka.consumer.Consumer;
import kafka.consumer.ConsumerConfig;
import kafka.consumer.KafkaStream;
import kafka.javaapi.consumer.ConsumerConnector;

/**
 * A wraper of Kafka consumer
 * 
 * @author Hongwei Wang
 * 
 */
public class KafkaConsumer {
	private ConsumerConnector consumer;
	private static Logger logger = Logger.getLogger(KafkaConsumer.class);

	/**
	 * Constructor, all consumers have the same group id
	 * 
	 * @param zookeeper
	 */
	public KafkaConsumer(String zookeeper) {
		this(zookeeper, "groupID");
	}

	/**
	 * you can specify the group id the consumer belonging to
	 * 
	 * @param zookeeper
	 * @param groupId
	 */
	public KafkaConsumer(String zookeeper, String groupId) {

		Properties props = new Properties();

		// the host and port of zookeeper server: localhost:2181
		props.put("zookeeper.connect", zookeeper);
		// all consumers belong to the same group
		logger.info("Groupid is " + groupId);
		props.put("group.id", groupId);
		props.put("zookeeper.session.timeout.ms", "400000");
		props.put("zookeeper.sync.time.ms", "40000");
		props.put("rebalance.backoff.ms", "40000");
		ConsumerConfig config = new ConsumerConfig(props);
		consumer = (ConsumerConnector) Consumer.createJavaConsumerConnector(config);
	}

	/**
	 * Consume message stream
	 * 
	 * @param topicCount
	 * @return
	 */
	public Map<String, List<KafkaStream<byte[], byte[]>>> consume(Map<String, Integer> topicCount) {
		Map<String, List<KafkaStream<byte[], byte[]>>> consumerStreams = consumer.createMessageStreams(topicCount);
		return consumerStreams;
	}

	public KafkaStream<byte[], byte[]> consumeTopic(String topic) {
		Map<String, Integer> topicCount = new HashMap<String, Integer>();
		topicCount.put(topic, 1);
		Map<String, List<KafkaStream<byte[], byte[]>>> consumerStreams = consumer.createMessageStreams(topicCount);

		return consumerStreams.get(topic).get(0);
	}

	/**
	 * Close consumer
	 */
	public void close() {
		consumer.shutdown();
	}
}
